import torch
import os.path as osp
import numpy as np
from torch.utils.data import Dataset

from joblib import Memory
from sklearn.datasets import load_svmlight_file

mem = Memory('./mycache')


@mem.cache
def get_data(file):
    data = load_svmlight_file(file)
    return data[0].todense(), data[1]


class ChemDataset(Dataset):

    def __init__(self, data_dir, train=True, local_test=False, test_only=False,seed=None, ratio=0.3,
                 subsample=10000):

        pos_data, _ = get_data(osp.join(data_dir, 'data_train.txt'))
        pos_data = pos_data[:521] # only use positive data

        self.ratio = ratio
        self.feat_dim = 1024
        self.train = train

        self.neg_data = None

        if not seed:
            self.seed = seed
            np.random.seed(seed)

        N_p = len(pos_data)

        if train:
            self.data = pos_data
            self.labels = np.ones((N_p,))
        elif test_only:
            mix_data, _ = get_data(osp.join(data_dir, 'data_predict.txt'))
            self.labels = np.zeros(len(mix_data),)
            self.data = mix_data
        else:
            if local_test:
                mix_data, _ = get_data(osp.join(data_dir,'data_valid.txt'))
                self.labels = np.concatenate((np.ones(179, ), np.zeros(len(mix_data) - 179, )))
            else:
                mix_data, _ = get_data(osp.join(data_dir,'data_valid.txt'))
                pos_data = mix_data[:179,:]

                neg_data, _ = get_data(osp.join(data_dir, 'data_predict.txt'))
                neg_data = mix_data[:3000]
                mix_data = np.concatenate((pos_data, neg_data), axis=0)

                self.labels = np.concatenate((np.ones(179,), np.zeros(3000, )))
            self.data = mix_data


    def __getitem__(self, index):
        return torch.from_numpy(self.data[index]).float().squeeze(), self.labels[index]

    def parse(self, feat):
        data = np.zeros(self.feat_dim)
        feats = [int(x.split(':')[0]) - 1 for x in feat.split(' ')[2:]]
        data[np.array(feats)] = 1
        return torch.from_numpy(data).float().squeeze()

    def __len__(self):

        return self.data.shape[0]


    def get_batch(self, batchsize):
        kp = np.random.permutation(len(self))
        if self.train:
            batch = torch.from_numpy(self.data[kp[:batchsize]]).float().squeeze()
            labels = torch.from_numpy(self.labels[kp[:batchsize]]).long()
            return batch, labels
        else:
            raise NotImplementedError("get_batch method in ChemDataset only support train mode")


if __name__ == '__main__':
    ds = ChemDataset('../data/', train=True, seed=1213)
    x, y = ds.get_batch(4)
    print(x.size(), y.size())
